{
 "nbformat": 4,
 "nbformat_minor": 0,
 "metadata": {
  "colab": {
   "name": "GPT",
   "provenance": [],
   "collapsed_sections": []
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "accelerator": "GPU"
 },
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AYV_dMVDxyc2"
   },
   "source": [
    "[![Github](https://img.shields.io/github/stars/labmlai/annotated_deep_learning_paper_implementations?style=social)](https://github.com/labmlai/annotated_deep_learning_paper_implementations)\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/transformers/gpt/experiment.ipynb)                    \n",
    "\n",
    "## Training a model with GPT architecture\n",
    "\n",
    "This is an experiment training Tiny Shakespeare dataset with GPT architecture model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AahG_i2y5tY9"
   },
   "source": [
    "Install the `labml-nn` package"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "ZCzmCrAIVg0L",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "9c544df4-3fc7-4152-b50d-07919dbfb9de"
   },
   "source": [
    "!pip install labml-nn"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SE2VUQ6L5zxI"
   },
   "source": [
    "Imports"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "0hJXx_g0wS2C"
   },
   "source": [
    "from labml import experiment\n",
    "from labml_nn.transformers.gpt import Configs"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Lpggo0wM6qb-"
   },
   "source": [
    "Create an experiment"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "bFcr9k-l4cAg"
   },
   "source": [
    "experiment.create(name=\"gpt\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-OnHLi626tJt"
   },
   "source": [
    "Initialize [GPT configurations](https://nn.labml.ai/transformers/gpt/)"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "Piz0c5f44hRo"
   },
   "source": [
    "conf = Configs()"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wwMzCqpD6vkL"
   },
   "source": [
    "Set experiment configurations and assign a configurations dictionary to override configurations"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 17
    },
    "id": "e6hmQhTw4nks",
    "outputId": "018b39d7-7d84-4651-8b33-80de77d58ace"
   },
   "source": [
    "experiment.configs(conf, {\n",
    "    # Use character level tokenizer\n",
    "    'tokenizer': 'character',\n",
    "    # Prompt separator is blank\n",
    "    'prompt_separator': '',\n",
    "    # Starting prompt for sampling\n",
    "    'prompt': 'It is ',\n",
    "    # Use Tiny Shakespeare dataset\n",
    "    'text': 'tiny_shakespeare',\n",
    "\n",
    "    # Use a context size of $128$\n",
    "    'seq_len': 128,\n",
    "    # Train for $32$ epochs\n",
    "    'epochs': 32,\n",
    "    # Batch size $128$\n",
    "    'batch_size': 128,\n",
    "    # Switch between training and validation for $10$ times\n",
    "    # per epoch\n",
    "    'inner_iterations': 10,\n",
    "\n",
    "    # Transformer configurations\n",
    "    'transformer.d_model': 512,\n",
    "    'transformer.ffn.d_ff': 2048,\n",
    "    'transformer.n_heads': 8,\n",
    "    'transformer.n_layers': 6\n",
    "})"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EvI7MtgJ61w5"
   },
   "source": [
    "Set PyTorch models for loading and saving"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 459
    },
    "id": "GDlt7dp-5ALt",
    "outputId": "e7768185-4a3c-4fec-f688-73c0463db015"
   },
   "source": [
    "experiment.add_pytorch_models({'model': conf.model})"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KJZRf8527GxL"
   },
   "source": [
    "Start the experiment and run the training loop."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "aIAWo7Fw5DR8",
    "outputId": "d4ca5ae0-6a97-439e-a318-010cb3f288cd"
   },
   "source": [
    "# Start the experiment\n",
    "with experiment.start():\n",
    "    conf.run()"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "oBXXlP2b7XZO"
   },
   "source": [
    ""
   ],
   "outputs": [],
   "execution_count": null
  }
 ]
}
